#!/usr/bin/env python3
"""
MCP 客户端 - 将本地MCP服务连接到小智AI
"""

import asyncio
import json
import logging
import subprocess
import websockets
import argparse
import os
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
import signal
import sys
from dotenv import load_dotenv

load_dotenv()
# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

@dataclass
class MCPService:
    name: str
    command: List[str]
    process: Optional[subprocess.Popen] = None
    tools: Dict[str, Any] = None
    initialized: bool = False

class MCPClient:
    def __init__(self, endpoint_url: str):
        self.xiaozhi_url = endpoint_url
        self.services: Dict[str, MCPService] = {}
        self.all_tools: Dict[str, Dict] = {}
        self.request_id = 0
        self.websocket = None
        self.running = True
        
    def get_next_id(self) -> int:
        self.request_id += 1
        return self.request_id
    
    async def add_service(self, name: str, command: List[str], env: Dict[str, str] = None) -> bool:
        """添加一个MCP服务"""
        try:
            logger.info(f"启动服务: {name}")
            
            # 获取当前环境变量
            process_env = os.environ.copy()
            
            # 合并自定义环境变量
            if env:
                process_env.update(env)
            
            # 启动子进程
            process = subprocess.Popen(
                command,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=0,
                env=process_env
            )
            
            service = MCPService(
                name=name,
                command=command,
                process=process,
                tools={}
            )
            
            self.services[name] = service
            
            # 初始化服务
            await self.initialize_service(name)
            
            # 获取工具列表
            await self.load_service_tools(name)
            
            logger.info(f"服务 {name} 启动成功")
            return True
            
        except Exception as e:
            logger.error(f"启动服务 {name} 失败: {e}")
            return False
    
    async def initialize_service(self, service_name: str):
        """初始化MCP服务"""
        service = self.services.get(service_name)
        if not service or not service.process:
            return
        
        init_request = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {
                    "name": "xiaozhi-mcp-bridge",
                    "version": "1.0.0"
                }
            }
        }
        
        try:
            service.process.stdin.write(json.dumps(init_request) + '\n')
            service.process.stdin.flush()
            

            # 读取响应
            response_line = service.process.stdout.readline()
            logger.info(f"初始化响应: {response_line}")
            # 等待初始化响应
            # await asyncio.sleep(1)
            service.initialized = True
            
        except Exception as e:
            logger.error(f"初始化服务 {service_name} 失败: {e}")
    
    async def load_service_tools(self, service_name: str):
        """加载服务的工具列表"""
        service = self.services.get(service_name)
        if not service or not service.process:
            return
        
        tools_request = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": "tools/list",
            "params": {}
        }
        
        try:
            service.process.stdin.write(json.dumps(tools_request) + '\n')
            service.process.stdin.flush()
            
            # 读取响应
            response_line = service.process.stdout.readline()
            logger.info(f"工具列表响应: {response_line}")
            if response_line:
                response = json.loads(response_line.strip())
                if response.get('result') and response['result'].get('tools'):
                    tools = response['result']['tools']
                    for tool in tools:
                        tool_name = tool['name']
                        # 添加服务名前缀避免冲突
                        prefixed_name = f"{service_name}_{tool_name}"
                        self.all_tools[prefixed_name] = {
                            **tool,
                            'name': prefixed_name,
                            'original_name': tool_name,
                            'service_name': service_name,
                            'description': f"[{service_name}] {tool.get('description', '')}"
                        }
                    
                    logger.info(f"从服务 {service_name} 加载了 {len(tools)} 个工具")
                    
        except Exception as e:
            logger.error(f"加载服务 {service_name} 工具失败: {e}")
    
    async def call_tool(self, tool_name: str, arguments: Dict = None) -> Dict:
        """调用工具"""
        if arguments is None:
            arguments = {}
            
        tool_info = self.all_tools.get(tool_name)
        if not tool_info:
            raise Exception(f"工具 {tool_name} 不存在")
        
        service_name = tool_info['service_name']
        original_tool_name = tool_info['original_name']
        service = self.services.get(service_name)
        
        if not service or not service.process:
            raise Exception(f"服务 {service_name} 不可用")
        
        request_id = self.get_next_id()
        call_request = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "tools/call",
            "params": {
                "name": original_tool_name,
                "arguments": arguments
            }
        }
        
        try:
            service.process.stdin.write(json.dumps(call_request) + '\n')
            service.process.stdin.flush()
            
            # 读取响应
            response_line = service.process.stdout.readline()
            if response_line:
                response = json.loads(response_line.strip())
                if response.get('error'):
                    raise Exception(response['error'].get('message', '工具调用失败'))
                return response.get('result', {})
            else:
                raise Exception("工具调用无响应")
                
        except Exception as e:
            logger.error(f"调用工具 {tool_name} 失败: {e}")
            raise
    
    async def handle_xiaozhi_message(self, message: Dict) -> Dict:
        """处理来自小智AI的消息"""
        method = message.get('method')
        params = message.get('params', {})
        msg_id = message.get('id')
        
        try:
            if method == 'tools/list':
                # 返回所有可用工具
                tools_list = list(self.all_tools.values())
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": {
                        "tools": tools_list
                    }
                }
            
            elif method == 'tools/call':
                # 调用工具
                tool_name = params.get('name')
                arguments = params.get('arguments', {})
                
                result = await self.call_tool(tool_name, arguments)
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": result
                }
            elif method == 'initialize':
                """处理初始化请求"""
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {
                            "tools": {}
                        },
                        "serverInfo": {
                            "name": "xiaozhi-mcp-bridge",
                            "version": "1.0.0"
                        }
                    }
                }
            elif method == 'notifications/initialized':
                """处理初始化通知"""
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": {}
                }
            elif method == 'ping':
                """处理ping请求"""
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": {}
                }
            else:
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "error": {
                        "code": -32601,
                        "message": f"未知方法: {method}"
                    }
                }
                
        except Exception as e:
            logger.error(f"处理消息失败: {e}")
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {
                    "code": -32000,
                    "message": str(e)
                }
            }
    
    async def connect_to_xiaozhi(self):
        """连接到小智AI"""
        logger.info(f"连接到小智AI: {self.xiaozhi_url}")
        
        try:
            async with websockets.connect(self.xiaozhi_url) as websocket:
                self.websocket = websocket
                logger.info("已连接到小智AI")
                
                # 发送初始化信息
                init_info = {
                    "jsonrpc": "2.0",
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {
                            "tools": {}
                        },
                        "serverInfo": {
                            "name": "local-mcp-client",
                            "version": "1.0.0"
                        }
                    }
                }
                await websocket.send(json.dumps(init_info))
                
                # 消息循环
                async for message in websocket:
                    if self.running == False:
                        break
                    try:
                        data = json.loads(message)
                        logger.debug(f"收到消息: {data}")
                        
                        response = await self.handle_xiaozhi_message(data)
                        if response:
                            await websocket.send(json.dumps(response))
                            logger.debug(f"发送响应: {response}")
                            
                    except json.JSONDecodeError:
                        logger.error(f"无效的JSON消息: {message}")
                    except Exception as e:
                        logger.error(f"处理消息时出错: {e}")
                        
        except websockets.exceptions.ConnectionClosed:
            logger.info("与小智AI的连接已关闭")
        except Exception as e:
            logger.error(f"连接小智AI失败: {e}")
            logger.info("5秒后重新连接...")
            await asyncio.sleep(5)
    
    async def start(self):
        """启动客户端"""
        logger.info("启动MCP客户端")
        
        # 设置信号处理
        def signal_handler(signum, frame):
            logger.info("收到退出信号，正在关闭...")
            self.running = False
            
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
        
        try:
            # 连接到小智AI
            while self.running:
                try:
                    await self.connect_to_xiaozhi()
                except Exception as e:
                    logger.error(f"连接失败: {e}")
                    if self.running:
                        logger.info("5秒后重新连接...")
                        await asyncio.sleep(5)
                    
        finally:
            await self.cleanup()
    
    async def cleanup(self):
        """清理资源"""
        logger.info("清理资源...")
        
        for service in self.services.values():
            if service.process:
                try:
                    service.process.terminate()
                    service.process.wait(timeout=5)
                except:
                    service.process.kill()
        
        if self.websocket:
            await self.websocket.close()
def load_config(config_file="config.json"):
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        return None
    
async def main():
    parser = argparse.ArgumentParser(description='MCP客户端 - 连接小智AI')
    # parser.add_argument('--token', help='小智AI的token')
    parser.add_argument('--debug', action='store_true', help='启用调试日志')
    
    args = parser.parse_args()
    
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
    
    config = load_config()

    endpoint_url = os.environ.get('MCP_ENDPOINT')
    if not endpoint_url:
        logger.error("Please set the `MCP_ENDPOINT` environment variable")
        sys.exit(1)
    client = MCPClient(endpoint_url)
    
    # 在main函数中使用
    if config:
        for service in config['services']:
            if service.get('enabled', True):
                success = await client.add_service(service['name'], service['command'], service.get('env', {}))
                if success:
                    # 等待一下避免同时启动太多进程
                    await asyncio.sleep(1)

    
    # 显示已加载的工具
    if client.all_tools:
        logger.info(f"已加载 {len(client.all_tools)} 个工具:")
        for tool_name, tool_info in client.all_tools.items():
            logger.info(f"  - {tool_name}: {tool_info.get('description', '')}")
    else:
        logger.warning("没有加载任何工具")
    
    # 启动客户端
    await client.start()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序运行出错: {e}")
        sys.exit(1)